#include <rclcpp/rclcpp.hpp>
#include <geometry_msgs/msg/twist.hpp>
#include <termios.h>
#include <unistd.h>
#include <cstdio>

// 获取键盘输入（非阻塞）
char getKey()
{
    struct termios oldt, newt;
    char ch = 0;

    tcgetattr(STDIN_FILENO, &oldt);
    newt = oldt;
    newt.c_lflag &= ~(ICANON | ECHO);
    tcsetattr(STDIN_FILENO, TCSANOW, &newt);

    fd_set set;
    struct timeval timeout;
    timeout.tv_sec = 0;
    timeout.tv_usec = 0;

    FD_ZERO(&set);
    FD_SET(STDIN_FILENO, &set);

    if (select(STDIN_FILENO + 1, &set, NULL, NULL, &timeout) > 0)
    {
        read(STDIN_FILENO, &ch, 1);
    }

    tcsetattr(STDIN_FILENO, TCSANOW, &oldt);
    return ch;
}

class WheelSpeedControl : public rclcpp::Node
{
public:
    // 创建一个Publisher，发布名为/cmd_vel的话题，消息类型为geometry_msgs::msg::Twist
    WheelSpeedControl() : Node("wheel_speed_control")
    {
        cmd_vel_pub_ = create_publisher<geometry_msgs::msg::Twist>("/cmd_vel", 10);
        // 定时器用于定期调用keyControl函数读取键盘输入
        timer_ = create_wall_timer(
            std::chrono::milliseconds(100),
            [this]()
            { this->keyControl(); });

        // 配置参数
        max_linear_speed_ = 0.5;  // 最大线速度 (m/s)
        max_angular_speed_ = 0.5; // 最大角速度 (rad/s)
        wheel_separation_ = 0.36; // 轮间距 (m)，与URDF中一致
        wheel_radius_ = 0.09;     // 轮子半径 (m)，直径0.18m，所以半径0.09m

        RCLCPP_INFO(get_logger(), "Wheel speed control ready:");
        RCLCPP_INFO(get_logger(), "  ↑: Forward");
        RCLCPP_INFO(get_logger(), "  ↓: Backward");
        RCLCPP_INFO(get_logger(), "  ←: Turn left");
        RCLCPP_INFO(get_logger(), "  →: Turn right");
        RCLCPP_INFO(get_logger(), "  Space: Stop");
        RCLCPP_INFO(get_logger(), "  q: Quit");
    }

private:
    // 将归一化轮速转换为Twist消息
    geometry_msgs::msg::Twist calculateTwist(double left_normalized, double right_normalized)
    {
        auto twist = geometry_msgs::msg::Twist();
        
        // 计算实际轮速 (m/s)
        double left_speed = left_normalized * max_linear_speed_;
        double right_speed = right_normalized * max_linear_speed_;
        
        // 根据差分驱动模型计算线速度和角速度
        // 线速度 = (左轮速度 + 右轮速度) / 2
        twist.linear.x = (left_speed + right_speed) / 2.0;
        
        // 角速度 = (右轮速度 - 左轮速度) / 轮间距
        twist.angular.z = (right_speed - left_speed) / wheel_separation_;
        
        return twist;
    }

    void keyControl()
    {
        char key = getKey();
        double left_normalized = 0.0;
        double right_normalized = 0.0;
        bool publish = true;

        switch (key)
        {
        case 'A': // 上箭头: 前进
            left_normalized = 1.0;
            right_normalized = 1.0;
            RCLCPP_INFO(get_logger(), "Moving forward");
            break;
        case 'B': // 下箭头: 后退
            left_normalized = -1.0;
            right_normalized = -1.0;
            RCLCPP_INFO(get_logger(), "Moving backward");
            break;
        case 'C': // 右箭头: 右转
            left_normalized = 0.5;
            right_normalized = -0.5;
            RCLCPP_INFO(get_logger(), "Turning right");
            break;
        case 'D': // 左箭头: 左转
            left_normalized = -0.5;
            right_normalized = 0.5;
            RCLCPP_INFO(get_logger(), "Turning left");
            break;
        case ' ': // 空格: 停止
            RCLCPP_INFO(get_logger(), "Stopping");
            break;
        case 'q': // 退出
            RCLCPP_INFO(get_logger(), "Exiting...");
            rclcpp::shutdown();
            exit(0);
        default:
            // 保持当前速度
            publish = false;
            if (key != 0)
            {
                RCLCPP_INFO(get_logger(), "Invalid key: %c", key);
            }
            break;
        }

        if (publish)
        {
            auto twist = calculateTwist(left_normalized, right_normalized);
            cmd_vel_pub_->publish(twist);
        }
    }

    rclcpp::Publisher<geometry_msgs::msg::Twist>::SharedPtr cmd_vel_pub_;
    rclcpp::TimerBase::SharedPtr timer_;
    
    // 控制参数
    double max_linear_speed_;
    double max_angular_speed_;
    double wheel_separation_;
    double wheel_radius_;
};

int main(int argc, char **argv)
{
    rclcpp::init(argc, argv);
    auto node = std::make_shared<WheelSpeedControl>();
    rclcpp::spin(node);
    rclcpp::shutdown();
    return 0;
}